#!/usr/bin/python3

import os
from jinja2 import Environment, PackageLoader

classes = [
    {
        "name"  : "Assert",
        "base"  : "stmtType",
        "id"    : "JJTASSERT_STMT",
        "params": [
            ("exprTypePtr", "test"),
            ("exprTypePtr", "msg")
        ]
    },
    {
        "name"  : "Assign",
        "base"  : "stmtType",
        "id"    : "-1",
        "params": [
            ("ExprTypeList", "targets"),
            ("exprTypePtr", "value")
        ],
    },
    {
        "name"  : "Attribute",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "value"),
            ("NameTokPtr", "attr"),
            ("int", "ctx")
        ],
        "addinc" : "exprContextType"
    },
    {
        "name"  : "AugAssign",
        "base"  : "stmtType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "target"),
            ("int", "op"),
            ("exprTypePtr", "value")
        ],
        "addinc": "operatorType"
    },
    {
        "name"  : "BinOp",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "left"),
            ("int", "op"),
            ("exprTypePtr", "right")
        ],
        "addinc": "operatorType"
    },
    {
        "name"  : "BoolOp",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("int", "op"),
            ("ExprTypeList", "values")
        ],
        "addinc": "boolopType"
    },
    {
        "name"  : "Break",
        "base"  : "stmtType",
        "id"    : "JJTBREAK_STMT",
        "params": [
        ],
    },
    {
        "name"  : "Call",
        "base"  : "exprType",
        "id"    : "JJTCALL_OP",
        "params": [
            ("exprTypePtr", "func"),
            ("ExprTypeList", "args"),
            ("KeywordTypeList", "keywords"),
            ("exprTypePtr", "starargs"),
            ("exprTypePtr", "kwargs"),
        ],
    },
    {
        "name"  : "ClassDef",
        "base"  : "stmtType",
        "id"    : "JJTCLASSDEF",
        "params": [
            ("NameTokPtr", "name"),
            ("ExprTypeList", "bases"),
            ("StmtTypeList", "body"),
            ("DecoratorsTypeList", "decs"),
            ("KeywordTypeList", "keywords"),
            ("exprTypePtr", "starargs"),
            ("exprTypePtr", "kwargs"),
        ],
    },
    {
        "name"  : "Compare",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "left"),
            ("IntList_", "ops"),
            ("ExprTypeList", "comparators")
        ],
        "addinc": ["cmpopType", "intList"]
    },
    {
        "name"  : "Comprehension",
        "base"  : "comprehensionType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "target"),
            ("exprTypePtr", "iter"),
            ("ExprTypeList", "ifs")
        ],
        "needList" : True
    },
    {
        "name"  : "Continue",
        "base"  : "stmtType",
        "id"    : "JJTCONTINUE_STMT",
        "params": [
        ],
    },
    {
        "name"  : "Delete",
        "base"  : "stmtType",
        "id"    : "JJTDEL_STMT",
        "params": [
            ("ExprTypeList", "target"),
        ],
    },
    {
        "name"  : "Dict",
        "base"  : "exprType",
        "id"    : "JJTDICTIONARY",
        "params": [
            ("ExprTypeList", "keys"),
            ("ExprTypeList", "values"),
        ],
    },
    {
        "name"  : "DictComp",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "key"),
            ("exprTypePtr", "value"),
            ("ComprehensionTypeList", "generators"),
        ],
    },
    {
        "name"  : "Ellipsis",
        "base"  : "sliceType",
        "id"    : "-1",
        "params": [
        ],
    },
    {
        "name"  : "Expr",
        "base"  : "stmtType",
        "id"    : "JJTEXPR_STMT",
        "params": [
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "Expression",
        "base"  : "modType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "body"),
        ],
    },
    {
        "name"  : "ExtSlice",
        "base"  : "sliceType",
        "id"    : "-1",
        "params": [
            ("SliceTypeList", "dims"),
        ],
    },
    {
        "name"  : "For",
        "base"  : "stmtType",
        "id"    : "JJTFOR_STMT",
        "params": [
            ("exprTypePtr", "target"),
            ("exprTypePtr", "iter"),
            ("StmtTypeList", "body"),
            ("suiteTypePtr", "orelse"),
        ],
    },
    {
        "name"  : "FunctionDef",
        "base"  : "stmtType",
        "id"    : "JJTFUNCDEF",
        "params": [
            ("NameTokPtr", "name"),
            ("argumentsTypePtr", "args"),
            ("StmtTypeList", "body"),
            ("DecoratorsTypeList", "decs"),
            ("exprTypePtr", "returns"),
        ],
    },
    {
        "name"  : "GeneratorExp",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "elt"),
            ("ComprehensionTypeList", "generators"),
        ],
    },
    {
        "name"  : "Global",
        "base"  : "stmtType",
        "id"    : "JJTGLOBAL_STMT",
        "params": [
            ("NameTokList", "names"),
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "If",
        "base"  : "stmtType",
        "id"    : "JJTIF_STMT",
        "params": [
            ("exprTypePtr", "test"),
            ("StmtTypeList", "body"),
            ("suiteTypePtr", "orelse"),
        ],
    },
    {
        "name"  : "IfExp",
        "base"  : "exprType",
        "id"    : "JJTIF_EXP",
        "params": [
            ("exprTypePtr", "test"),
            ("exprTypePtr", "body"),
            ("exprTypePtr", "orelse"),
        ],
    },
    {
        "name"  : "Import",
        "base"  : "stmtType",
        "id"    : "JJTIMPORT",
        "params": [
            ("AliasTypeList", "names"),
        ],
    },
    {
        "name"  : "ImportFrom",
        "base"  : "stmtType",
        "id"    : "JJTIMPORTFROM",
        "params": [
            ("NameTokPtr", "module"),
            ("AliasTypeList", "names"),
            ("int", "level"),
        ],
    },
    {
        "name"  : "Index",
        "base"  : "sliceType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "Interactive",
        "base"  : "modType",
        "id"    : "-1",
        "params": [
            ("StmtTypeList", "body"),
        ],
    },
    {
        "name"  : "Lambda",
        "base"  : "exprType",
        "id"    : "JJTLAMBDEF",
        "params": [
            ("argumentsTypePtr", "args"),
            ("exprTypePtr", "body"),
        ],
    },
    {
        "name"  : "List",
        "base"  : "exprType",
        "id"    : "JJTLIST",
        "params": [
            ("ExprTypeList", "elts"),
            ("int", "ctx"),
        ],
        "addinc" : "exprContextType"
    },
    {
        "name"  : "ListComp",
        "base"  : "exprType",
        "id"    : "JJTLIST",
        "params": [
            ("exprTypePtr", "elt"),
            ("ComprehensionTypeList", "generators"),
            ("int", "ctx"),
        ],
        "addinc": "compContextType"
    },
    {
        "name"  : "Module",
        "base"  : "modType",
        "id"    : "JJTFILE_INPUT",
        "params": [
            ("StmtTypeList", "body"),
        ],
    },
    {
        "name"  : "Name",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("QString", "id"),
            ("int", "ctx"),
            ("bool", "reserved"),
        ],
        "addinc" : "exprContextType"
    },
    {
        "name"  : "NameTok",
        "base"  : "NameTokType",
        "id"    : "-1",
        "params": [
            ("QString", "id"),
            ("int", "ctx"),
        ],
        "addinc": "nameContextType",
        "needList": True
    },
    {
        "name"  : "NameTokType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept": True
    },
    {
        "name"  : "NonLocal",
        "base"  : "stmtType",
        "id"    : "JJTNONLOCAL_STMT",
        "params": [
            ("NameTokList", "names"),
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "Num",
        "base"  : "exprType",
        "id"    : "JJTNUM",
        "params": [
            ("ObjectPtr", "n"),
            ("int", "type"),
            ("QString", "num"),
        ],
        "addinc": "numTypeType"
    },
    {
        "name"  : "Pass",
        "base"  : "stmtType",
        "id"    : "JJTPASS_STMT",
        "params": [
        ],
    },
    {
        "name"  : "Raise",
        "base"  : "stmtType",
        "id"    : "JJTRAISE_STMT",
        "params": [
            ("exprTypePtr", "type"),
            ("exprTypePtr", "inst"),
            ("exprTypePtr", "tback"),
            ("exprTypePtr", "cause"),
        ],
    },
    {
        "name"  : "Repr",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "Return",
        "base"  : "stmtType",
        "id"    : "JJTRETURN_STMT",
        "params": [
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "Set",
        "base"  : "exprType",
        "id"    : "JJTSET",
        "params": [
            ("ExprTypeList", "elts"),
        ],
    },
    {
        "name"  : "SetComp",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "elt"),
            ("ComprehensionTypeList", "generators"),
        ],
    },
    {
        "name"  : "Slice",
        "base"  : "sliceType",
        "id"    : "JJTSLICE",
        "params": [
            ("exprTypePtr", "lower"),
            ("exprTypePtr", "upper"),
            ("exprTypePtr", "step"),
        ],
    },
    {
        "name"  : "Starred",
        "base"  : "exprType",
        "id"    : "JJTSTAR_EXPR",
        "params": [
            ("exprTypePtr", "value"),
            ("int", "ctx"),
        ],
        "addinc" : "exprContextType"
    },
    {
        "name"  : "Str",
        "base"  : "exprType",
        "id"    : "JJTSTRING",
        "params": [
            ("QString", "s"),
            ("int", "type"),
            ("bool", "unicode"),
            ("bool", "raw"),
            ("bool", "binary"),
        ],
        "addinc" : "strTypeType"
    },
    {
        "name"  : "StrJoin",
        "base"  : "exprType",
        "id"    : "JJTSTRJOIN",
        "params": [
            ("ExprTypeList", "strs"),
        ],
    },
    {
        "name"  : "Subscript",
        "base"  : "exprType",
        "id"    : "JJTINDEX_OP",
        "params": [
            ("exprTypePtr", "value"),
            ("sliceTypePtr", "slice"),
            ("int", "ctx"),
        ],
        "addinc" : "exprContextType"
    },
    #{
    #    "name"  : "Suite",
    #    "base"  : "suiteType",
    #    "id"    : "-1",
    #    "params": [
    #        ("StmtTypeList", "body"),
    #    ],
    #},
    {
        "name"  : "TryExcept",
        "base"  : "stmtType",
        "id"    : "JJTBEGIN_TRY_STMT",
        "params": [
            ("StmtTypeList", "body"),
            ("ExcepthandlerTypeList", "handlers"),
            ("suiteTypePtr", "orelse"),
        ],
    },
    {
        "name"  : "TryFinally",
        "base"  : "stmtType",
        "id"    : "JJTTRYFINALLY_STMT",
        "params": [
            ("StmtTypeList", "body"),
            ("suiteTypePtr", "finalbody"),
        ],
    },
    {
        "name"  : "Tuple",
        "base"  : "exprType",
        "id"    : "JJTTUPLE",
        "params": [
            ("ExprTypeList", "elts"),
            ("int", "ctx"),
            ("bool", "endsWithComma"),
        ],
        "addinc" : "exprContextType"
    },
    {
        "name"  : "UnaryOp",
        "base"  : "exprType",
        "id"    : "-1",
        "params": [
            ("int", "op"),
            ("exprTypePtr", "operand"),
        ],
        "addinc": "unaryopType"
    },
    {
        "name"  : "While",
        "base"  : "stmtType",
        "id"    : "JJTWHILE_STMT",
        "params": [
            ("exprTypePtr", "test"),
            ("StmtTypeList", "body"),
            ("suiteTypePtr", "orelse"),
        ],
    },
    {
        "name"  : "With",
        "base"  : "stmtType",
        "id"    : "-1",
        "params": [
            ("WithItemTypeList", "with_item"),
            ("suiteTypePtr", "body"),
        ],
    },
    {
        "name"  : "WithItem",
        "base"  : "WithItemType",
        "id"    : "JJTWITH_ITEM",
        "params": [
            ("exprTypePtr", "context_expr"),
            ("exprTypePtr", "optional_vars"),
        ],
    },
    {
        "name"  : "WithItemType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "Yield",
        "base"  : "exprType",
        "id"    : "JJTYIELD_EXPR",
        "params": [
            ("exprTypePtr", "value"),
        ],
    },
    {
        "name"  : "aliasType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("NameTokPtr", "name"),
            ("NameTokPtr", "asname"),
        ],
        "needList": True,
        "accept": True
    },
    {
        "name"  : "argumentsType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("ExprTypeList", "args"),
            ("NameTokPtr", "vararg"),
            ("NameTokPtr", "kwarg"),
            ("ExprTypeList", "defaults"),
            ("ExprTypeList", "kwonlyargs"),
            ("ExprTypeList", "kw_defaults"),
            ("ExprTypeList", "annotation"),
            ("exprTypePtr", "varargannotation"),
            ("exprTypePtr", "kwargannotation"),
            ("ExprTypeList", "kwonlyargannotation"),
        ],
        "accept": True,
    },
    {
        "name"  : "commentType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("QString", "id"),
        ],
        "accept" : True
    },
    {
        "name"  : "comprehensionType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept": True,
    },
    {
        "name"  : "decoratorsType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "func"),
            ("ExprTypeList", "args"),
            ("KeywordTypeList", "keywords"),
            ("exprTypePtr", "starargs"),
            ("exprTypePtr", "kwargs"),
        ],
        "accept": True,
        "needList": True,
    },
    {
        "name"  : "excepthandlerType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("exprTypePtr", "type"),
            ("exprTypePtr", "name"),
            ("StmtTypeList", "body"),
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "exprType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "keywordType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("NameTokPtr", "arg"),
            ("exprTypePtr", "value"),
            ("bool", "afterstarargs"),
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "modType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "sliceType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "stmtType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
        ],
        "needList": True,
        "accept" : True
    },
    {
        "name"  : "suiteType",
        "base"  : "SimpleNode",
        "id"    : "-1",
        "params": [
            ("StmtTypeList", "body")
        ],
        "needList": True,
        "accept" : True
    },
]

def param(value):
    if value.startswith("Q") or value.endswith("List") or value.endswith("_") or value.endswith("Ptr"):
        return "const "+value+"&"
    return value
    
def normaltraverse(value):
    if value.endswith("Ptr") or value.endswith("*"):
        return True
    return False
    
def listtraverse(value):
    if value.endswith("List"):
        return True
    return False
    
def cap(value):
    return value[:1].upper()+value[1:]

def norm(value):
    return value+"_" if value in ("type") else value

def filterIncludes(params, base):
    ret = []
    incs = {}
    #ret.append(base)
    for param in params:
        incs[param[0]] = 1
    for inc in incs.keys():
        if inc.endswith("*") and inc[:-1] != base:
            ret.append(inc[:-1])
        if inc.endswith("Ptr") and inc[:-3] != base:
            ret.append(inc[:-3])
        if inc.endswith("List") and inc[:-4] != base:
            global classes
            found = False
            for cls in classes:
                if cls["name"].lower() == inc[:-4].lower():
                    ret.append(cls["name"])
                    found = True
            if not found:
                ret.append(inc[:-4])
    return ret
    
env = Environment(
    loader=PackageLoader('__main__', "."),
    line_statement_prefix="//"
)

env.filters['param'] = param
env.filters['normaltraverse'] = normaltraverse
env.filters['listtraverse'] = listtraverse
env.filters['cap'] = cap
env.filters['norm'] = norm



tmpl = env.get_template("header.tmpl")
cpptmpl = env.get_template("source.tmpl")

for cls in classes:
    hs = tmpl.render(
        className = cls["name"],
        classBase = cls["base"],
        nodeId    = cls["id"],
        params    = cls["params"],
        includes  = filterIncludes(cls["params"], cls["base"]),
        needList  = cls.get("needList", False),
        accept    = cls.get("accept", False),
        addinc    = cls.get("addinc", []) if isinstance(cls.get("addinc", []), list) else [cls.get("addinc", "")]
    )
    ss = cpptmpl.render(
        className = cls["name"],
        classBase = cls["base"],
        nodeId    = cls["id"],
        params    = cls["params"],
        includes  = filterIncludes(cls["params"], cls["base"]),
        needList  = cls.get("needList", False),
        accept    = cls.get("accept", False),
        addinc    = cls.get("addinc", []) if isinstance(cls.get("addinc", []), list) else [cls.get("addinc", "")]
    )
    
    with open("../"+cls["name"]+".h", "w") as out:
        out.write(hs)
    with open("../"+cls["name"]+".cpp", "w") as out:
        out.write(ss)

tmpl = env.get_template("visitor.tmpl")
visit = []
for cls in classes:
    if not cls.get("accept", False):
        visit.append(cls)
    
with open("../VisitorBase.h", "w") as out:
    out.write(tmpl.render(visit = visit))
    
tmpl = env.get_template("ivisitor.tmpl")
with open("../../IVisitor.h", "w") as out:
    out.write(tmpl.render(visit = visit))

tmpl = env.get_template("asttypes.tmpl")
with open("../../AstTypes.h", "w") as out:
    out.write(tmpl.render(visit = classes))
    
tmpl = env.get_template("asttypes-cpp.tmpl")
with open("../../AstTypes.cpp", "w") as out:
    out.write(tmpl.render(visit = classes))
